在昇腾 NPU 上运行 RLVR 流水线
最后更新:2026/04/28。
本文档提供在华为昇腾 NPU 上运行 RLVR(Reinforcement Learning with Verifiable Rewards)流水线的端到端指南,涵盖环境准备、数据准备、模型下载、配置编写、训练启动、监控与评估,以及从 checkpoint 恢复训练。
工作流概览
从零开始在 NPU 上运行 RLVR 任务包含以下步骤:
1. 环境准备 → 2. 数据准备 → 3. 模型准备 → 4. 编写配置 → 5. 启动训练 → 6. 监控与评估 → 7. 从 Checkpoint 恢复
步骤 1:环境准备
1.1 硬件与驱动前置条件
请确保硬件和宿主机驱动已经准备就绪:
| 项目 | 要求 |
|---|---|
| 硬件 | Atlas 900 A2 PODc(Ascend 910B1)或 Atlas 900 A3 PODc(Ascend 910_9391) |
| 宿主机 OS | Ubuntu 22.04 |
| CANN | 9.0.0 |
| Ascend NPU 驱动 | 已在宿主机安装(npu-smi info 能看到设备) |
| Docker | >= 20.10 |
1.2 获取 Docker 镜像
请使用与硬件匹配的预构建昇腾镜像。官方 ROLL NPU 镜像标签可在 https://quay.io/repository/ascend/roll?tab=tags 查看。容器启动细节参见 Ascend NPU Docker 使用指南。
# A2 硬件
docker pull quay.io/ascend/roll:main-a2
docker tag quay.io/ascend/roll:main-a2 roll:ascend-a2
# A3 硬件
docker pull quay.io/ascend/roll:main-a3
docker tag quay.io/ascend/roll:main-a3 roll:ascend-a3
当前仓库提供 docker/Dockerfile.A2 和 docker/Dockerfile.A3,用于构建自定义镜像。如果你维护自定义镜像,请确保依赖版本与预构建镜像保持一致。
1.3 启动容器
docker run -dit \
--name roll_npu \
--ulimit nofile=65536:65536 \
--device /dev/davinci0 \
--device /dev/davinci1 \
--device /dev/davinci2 \
--device /dev/davinci3 \
--device /dev/davinci4 \
--device /dev/davinci5 \
--device /dev/davinci6 \
--device /dev/davinci7 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
-v /usr/local/Ascend/driver:/usr/local/Ascend/driver \
-v /usr/local/Ascend/add-ons:/usr/local/Ascend/add-ons \
-v /usr/local/dcmi:/usr/local/dcmi \
-v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
-v /etc/ascend_install.info:/etc/ascend_install.info \
-v /path/to/models:/data/models \
-v /path/to/data:/data \
--ipc=host \
--net=host \
roll:ascend-a3 \
/bin/bash
注意:
-v /path/to/models:/data/models和-v /path/to/data:/data分别挂载模型权重目录和训练数据目录。请根据实际环境调整路径。
1.4 验证环境
进入容器后,执行:
# 验证 NPU 可见性
npu-smi info
# 验证 CANN 环境是否已加载
env | grep -E "ASCEND|LD_LIBRARY_PATH|PATH"
# 验证 Python 包
python -c "import torch; import torch_npu; print(torch.npu.is_available())"
python -c "import vllm; print(f'vllm: {vllm.__version__}')"
python -c "import vllm_ascend; print(f'vllm_ascend available')"
如果以上验证均通过,说明环境已经准备就绪。环境变量的详细说明参见 NPU 环境配置指南。
步骤 2:数据准备
RLVR 流水线使用 JSONL 格式的数据文件。不同奖励领域需要不同的数据字段。
2.1 数据格式
通用字段(所有领域都必需)
| 字段 | 类型 | 是否必需 | 说明 |
|---|---|---|---|
id | string/int | 是 | 数据样本的唯一标识 |
messages 或 prompt | string | 是 | 输入 prompt;messages 是消息列表的 JSON 字符串 |
tag | string | 是 | 奖励领域标签,用于决定使用哪个 Reward Worker |
领域专属字段
| 领域 | tag 值 | 必需字段 | 说明 |
|---|---|---|---|
| 数学规则 | math_rule | ground_truth | 正确答案 |
| 代码沙箱 | code_sandbox(如 KodCode) | test_cases, case_type | 测试用例和类型(如 pytest) |
| LLM Judge | llm_judge(如 RLVR) | ground_truth | 参考答案或参考回复 |
| IFEval | ifeval | 无额外字段 | 基于规则的指令遵循评估 |
| CrossThinkQA | crossthinkqa | ground_truth | 跨学科推理答案 |
数据样例
数学领域(math_rule):
{
"id": "0",
"source": "gsm8k",
"difficulty": 0,
"prompt": "Solve the equation 3x + 5 = 14",
"messages": "[{\"role\": \"system\", \"content\": \"You are a math assistant.\"}, {\"role\": \"user\", \"content\": \"Solve the equation 3x + 5 = 14\"}]",
"ground_truth": "3",
"tag": "math_rule"
}
代码领域(code_sandbox):
{
"id": "5ea1ab",
"source": "codeforces",
"difficulty": "0",
"prompt": "Write a function that takes an array of distinct integers and returns all possible permutations.",
"messages": "[{\"role\": \"user\", \"content\": \"Write a function...\"}]",
"ground_truth": "[\"def permute(nums): ...\"]",
"case_type": "pytest",
"test_case_function": "",
"test_cases": "[{\"assert_code\": \"def test_permute(): ...\"}]",
"tag": "KodCode"
}
2.2 数据放置
将数据文件放到容器内的某个目录中(如 /data/),并在 actor_train.data_args 中指定路径:
actor_train:
data_args:
file_name:
- data/math_deepmath_deal.jsonl
- data/code_KodCode_data.jsonl
dataset_dir: data
2.3 验证数据
验证数据用于训练过程中的周期性评估。请在 validation 配置中指定:
validation:
data_args:
template: qwen2_5
file_name:
- data/math_benchmarks.jsonl
generating_args:
max_new_tokens: ${response_length}
top_p: 0.6
temperature: 0.6
num_return_sequences: 1
验证数据中的 tag 字段应与训练数据中的 tag 保持一致,这样才能按领域报告准确率。
步骤 3:模型准备
3.1 下载模型权重
RLVR 流水线需要以下模型:
| 模型 | 配置键 | 说明 |
|---|---|---|
| Actor / Reference 模型 | pretrain | 用于训练和推理的策略模型 |
| Reward 模型 | reward_pretrain | Reward Worker 中使用的模型(如数学规则奖励中的答案抽取) |
以 Qwen2.5-7B 为例:
# 使用 ModelScope 下载(推荐中国大陆用户使用)
pip install modelscope
modelscope download --model Qwen/Qwen2.5-7B --local_dir /data/models/Qwen2.5-7B
# 或使用 HuggingFace 下载
huggingface-cli download Qwen/Qwen2.5-7B --local-dir /data/models/Qwen2.5-7B
3.2 在配置中指定模型路径
pretrain: Qwen/Qwen2.5-7B # 从 ModelScope/HuggingFace 自动下载
# 或使用本地路径
# pretrain: /data/models/Qwen2.5-7B
reward_pretrain: Qwen/Qwen2.5-7B
提示: 如果容器内网络访问受限,请提前在宿主机下载模型,通过
-v挂载到容器中,并在配置中使用本地路径。
步骤 4:编写 NPU 配置
与 GPU 的关键差异
将 GPU RLVR 配置适配到 NPU 时,必须进行以下修改:
| 项目 | GPU | NPU |
|---|---|---|
| 训练后端 | Megatron 或 FSDP2 | 仅 FSDP2(NPU 不支持 Megatron) |
| 推理后端 | vLLM | vLLM-Ascend |
| Reference 模型策略 | megatron_infer | fsdp2_infer |
| 注意力实现 | flash_attn 或 fa2 | 通过 transformers 使用 fa2(不能使用 flash_attn 包) |
| 通信后端 | NCCL | HCCL |
| 设备可见性 | CUDA_VISIBLE_DEVICES | ASCEND_RT_VISIBLE_DEVICES |
| 分片配置 | FSDP2 或 Megatron 优化器分片 | FSDP2,7B+ 模型推荐 offload_policy: true |